#include "header.h"

static fftw_complex *a[6];
static fftw_complex *h[6];
static fftw_complex *zeta, *zeta_p;
static fftw_complex *beta, *beta_p;
static fftw_plan pf0, pb0;
static double *k, *v, overN, g, dt, dv;
static int N;
static vect rh[4], aux; 
static work_ptr tmp;

void err_msg(char* str)
{
  printf("%s\n",str);
  exit(1);
}


void init(int *argc, char **argv, params_ptr in, work_ptr wrk)
{
  char str[80], line[80], value[80], param[80];
  sprintf(str, "Usage:\n\t%s input", argv[0]);
  if (*argc != 2) err_msg(str);

  sprintf(str, "%s.cfg", argv[1]);
  FILE *fh = fopen(str,"r"); 
  FILE *fhlog = fopen("run.log","w"); 
  if (fh == NULL) err_msg("Cannot open file");
  while (fgets(line, 80, fh)!=NULL) {
    sscanf(line, "%s\t%s", param, value);
    if (strcmp(param,"#runname=") == 0) sprintf(in->runname,"%s", value);
    if (strcmp(param,"#resname=") == 0) sprintf(in->resname,"%s", value);
    if (strcmp(param,"#npoints=") == 0) in->n = atoi(value);
    if (strcmp(param,"#kcutoff=") == 0) in->k_cut = atoi(value);
    if (strcmp(param,"#cfl_con=") == 0) in->cfl = atof(value);
    if (strcmp(param,"#gravity=") == 0) in->g = atof(value);
    if (strcmp(param,"#stokesw=") == 0) in->stokes = atoi(value);
    if (strcmp(param,"#velocit=") == 0) in->c = atof(value);
    if (strcmp(param,"#fintime=") == 0) in->tmax = atof(value);
    if (strcmp(param,"#io_time=") == 0) in->t_skip = atof(value);
    if (strcmp(param,"#decay_p=") == 0) in->decay_p = atoi(value);
    if (strcmp(param,"#force_A=") == 0) in->force_A = atof(value);
    if (strcmp(param,"#k_min=") == 0) in->k_min = atoi(value);
    if (strcmp(param,"#k_max=") == 0) in->k_max = atoi(value);
  }
  fprintf(fhlog, "runname = %s\n", in->runname );
  fprintf(fhlog, "resname = %s\n", in->resname );
  fprintf(fhlog, "npoints = %d\n", in->n );
  fprintf(fhlog, "kcutoff = %d\n", in->k_cut );
  fprintf(fhlog, "cfl_con = %f\n", in->cfl );
  fprintf(fhlog, "gravity = %f\n", in->g );
  fprintf(fhlog, "stokesw = %d\n", in->stokes );
  fprintf(fhlog, "velocit = %f\n", in->c );
  fprintf(fhlog, "t_maxim = %f\n", in->tmax );
  fprintf(fhlog, "io_time = %f\n", in->t_skip );
  fclose(fh); fclose(fhlog);

  wrk->chi   = fftw_malloc(sizeof(fftw_complex)*(in->n));  
  wrk->rho   = fftw_malloc(sizeof(fftw_complex)*(in->n));
  wrk->Z   = fftw_malloc(sizeof(fftw_complex)*(in->n));
  wrk->Z_u = fftw_malloc(sizeof(fftw_complex)*(in->n));
  wrk->Phi = fftw_malloc(sizeof(fftw_complex)*(in->n));
  wrk->P_u = fftw_malloc(sizeof(fftw_complex)*(in->n));
  for (int j = 0; j < 4; j++) {
    rh[j].rho = fftw_malloc(sizeof(fftw_complex)*(in->n));
    rh[j].chi = fftw_malloc(sizeof(fftw_complex)*(in->n));
  }
  zeta   = fftw_malloc(sizeof(fftw_complex)*(in->n));
  zeta_p = fftw_malloc(sizeof(fftw_complex)*(in->n));
  beta   = fftw_malloc(sizeof(fftw_complex)*(in->n));
  beta_p = fftw_malloc(sizeof(fftw_complex)*(in->n));
  aux.rho = fftw_malloc(sizeof(fftw_complex)*(in->n));
  aux.chi = fftw_malloc(sizeof(fftw_complex)*(in->n));
  for (int j = 0; j < 6; j++) { 
    a[j] = fftw_malloc(sizeof(fftw_complex)*(in->n));
    h[j] = fftw_malloc(sizeof(fftw_complex)*(in->n));
  }
  v = fftw_malloc(sizeof(double)*(in->n));
  k = fftw_malloc(sizeof(double)*(in->n));

  wrk->alpha = 0.02*pi;
  for (int j = 0; j < in->n; j ++) {
    v[j] = 2*pi*(j+0.5)/(in->n);
    k[j] = j;
    if (j > (in->n)/2) k[j] = j - (in->n);
    wrk->rho[j] = 0.;/*   0.1*creal(theta(0.43, 0.01, v[j]+1.I*(wrk->alpha))) + 
	        0.05*creal(theta(0.56, 0.08, v[j]+1.I*(wrk->alpha)));*/
    wrk->chi[j] = 0;
  }
  pf0 = fftw_plan_dft_1d(in->n, a[0], a[1], FFTW_FORWARD, fmode);
  pb0 = fftw_plan_dft_1d(in->n, a[1], a[0], FFTW_BACKWARD, fmode);
  N = in->n; overN = 1./N; dv = 2*pi*overN;
  g = in->g; dt = (in->cfl)*2.*pi/N; 
  in->w_ptr = wrk; 
  in->num_iter = floor((in->tmax)/dt) + 1;
  dt = (in->tmax)/(in->num_iter); in->dt = dt;
  in->skip = floor((in->t_skip)/dt);
  srand(time(NULL));
}

void der(fftw_complex *in, fftw_complex *out) {
  memcpy(a[0], in, sizeof(fftw_complex)*N);
  fftw_execute(pf0);
  for (int j = 0; j < N ; j++) a[1][j] = 1.I*k[j]*overN*a[1][j];
  fftw_execute(pb0);
  memcpy(out, a[0], sizeof(fftw_complex)*N);  
}

void r_transform(fftw_complex *in, fftw_complex *out, double alpha) {
  memcpy(a[0], in, sizeof(fftw_complex)*N);
  fftw_execute(pf0);
  for (int j = 0; j < N ; j++) {
    a[1][j] = 1.I*ctanh(alpha*k[j])*overN*a[1][j];
    //printf("%e\t%e\n",creal(a[1][j]),cimag(a[1][j]));
  }
  fftw_execute(pb0);
  memcpy(out, a[0], sizeof(fftw_complex)*N);  
}

void t_transform(fftw_complex *in, fftw_complex *out, double alpha) {
  memcpy(a[0], in, sizeof(fftw_complex)*N);
  fftw_execute(pf0);
  for (int j = 0; j < N ; j++) a[1][j] = -1.I/(ctanh(alpha*k[j]))*overN*a[1][j];
  a[1][0] = 0;
  fftw_execute(pb0);
  memcpy(out, a[0], sizeof(fftw_complex)*N);  
}

void projector(fftw_complex *in, fftw_complex *out, double alpha) {
  memcpy(a[0], in, sizeof(fftw_complex)*N);
  fftw_execute(pf0);
  for (int j = 0; j < N ; j++) {
    a[1][j] = (1.-ctanh(alpha*k[j]))*overN*a[1][j];
  }
  fftw_execute(pb0);
  memcpy(out, a[0], sizeof(fftw_complex)*N);  
}

void iprojector(fftw_complex *in, fftw_complex *out, double alpha) {
  memcpy(a[0], in, sizeof(fftw_complex)*N);
  fftw_execute(pf0);
  for (int j = 1; j < N ; j++) {
    a[1][j] = (1. - 1./ctanh(alpha*k[j]) )*overN*a[1][j];
  }
  a[1][0] = overN*a[1][0];
  fftw_execute(pb0);
  memcpy(out, a[0], sizeof(fftw_complex)*N);  
}


void ffilter(params_ptr in) {
  tmp = in->w_ptr;
  double alfa = tmp->alpha;

  fft(tmp->rho, h[0]);
  fft(tmp->chi, h[1]);
  for (int j = 0; j < N; j++) {
    a[3][j] = -alfa*(in->force_A)*(alfa*k[j] - (in->k_min))*(alfa*k[j] - (in->k_max)); 
    if (creal(a[3][j]) < 0) a[3][j] = 0.;
    if (j > N/2) a[3][j] = 0.;
    a[3][j] = a[3][j]*cexp(2.I*pi*rand()/RAND_MAX);
    a[2][j] = a[2][j]*cexp(2.I*pi*rand()/RAND_MAX);
  }
  for (int j = 0; j < N/2; j++) {
    a[3][N-j-1] = conj(a[3][j+1]);
    a[2][N-j-1] = conj(a[2][j+1]);
  }
  a[2][N/2] = 0;
  a[3][N/2] = 0;

  for (int j = 0; j < N; j++) {
    a[4][j] = overN*cexp(-dt*pow(k[j]/(in->k_cut),in->decay_p));
   
    h[0][j] = h[0][j]*a[4][j];
    h[1][j] = (h[1][j] + dt*a[3][j])*a[4][j];
  }
  ifft(h[0], tmp->rho);
  ifft(h[1], tmp->chi);
  //write_out(a[4],"filter_spectrum.txt");
  //write_out(a[3],"forcing_spectrum.txt");
}

void fforcing(params_ptr in) {
  tmp = in->w_ptr;
  fft(tmp->chi, h[1]);
  for (int j = 0; j < N; j++) {
    a[4][j] = -(in->force_A)*
           (k[j] - (in->k_min))*(k[j] - (in->k_max));  //(tmp->alpha)*
    if (creal(a[4][j]) < 0) a[4][j] = 0.;
    if (j > N/2) a[4][j] = 0.;
  }
  write_out(a[4],"forcing_spectrum.txt");
  for (int j = 0; j < N; j++) {
    a[4][j]  = a[4][j]*cexp(2.I*pi*rand()/RAND_MAX);
    h[1][j]  = (h[1][j] + dt*a[4][j])*overN;
  }
  ifft(h[1], tmp->chi);
  //printf("%e\t%e\n",(tmp->alpha)*(in->k_min), (tmp->alpha)*(in->k_max));
  write_out(a[4], "forcing.txt"); 
}


void fft(fftw_complex *in, fftw_complex *out) {
  memcpy(a[0], in, N*sizeof(fftw_complex));
  fftw_execute(pf0);
  memcpy(out, a[1], N*sizeof(fftw_complex));
}

void ifft(fftw_complex *in, fftw_complex *out) {
  memcpy(a[1], in, N*sizeof(fftw_complex));
  fftw_execute(pb0);
  memcpy(out, a[0], N*sizeof(fftw_complex));
}

fftw_complex theta(fftw_complex r, fftw_complex d, fftw_complex w) {
   return -1.I*clog( (1+r*cexp(-1.I*(w-2*pi*d)))/(1+r*cexp(1.I*(w-2*pi*d)))  );
}

void hamiltonian(params_ptr in) {
  tmp = in->w_ptr;
  in->current.H = 0;
  projector(tmp->chi, beta, tmp->alpha);
  projector(tmp->rho, h[2], tmp->alpha);
  der(h[2], h[3]);
  for (int j = 0; j < N; j++) {
    zeta[j] = v[j] + 1.I*(tmp->alpha) + h[2][j];
    zeta_p[j] = 1 + h[3][j];
    h[4][j] = creal(beta[j]*zeta_p[j]);
  }
  fft(h[4],h[5]);
  for (int j = 1; j < N; j++) {
    (in->current.H) += pi*pow(cabs(h[5][j])*overN,2)*ctanh((tmp->alpha)*k[j])/k[j];
  }
  (in->current.H) += pi*pow(cabs(h[5][0])*overN,2)*(tmp->alpha);
  for (int j = 0; j < N; j++) {
    (in->current.H) += 0.5*dv*pow(cimag(zeta[j]),2)*creal(zeta_p[j]);
  }
}

void rhs(work_ptr in, int m) {
  projector(in->rho, h[2], in->alpha);
  projector(in->chi, beta, in->alpha);
  der(h[2], h[3]);
  der(beta, beta_p);
  for (int j = 0; j < N; j++) {
    zeta[j]   = v[j] + 1.I*(in->alpha) + h[2][j];
    zeta_p[j] = 1 + h[3][j];
    h[4][j] = cimag(-beta[j]/conj(zeta_p[j]));
    h[5][j] = 0.5*pow(cabs(beta[j]),2)+cimag(zeta[j]);
  } 
  der(h[5],h[5]);
  rh[m].alpha = 0.;
  for (int j = 0; j < N; j++) rh[m].alpha += overN*h[4][j];

  iprojector(h[4], h[0], in->alpha);
  projector(h[5],  h[1], in->alpha);
  for (int j = 0; j < N; j++) {
    rh[m].rho[j] = creal(1.I*zeta_p[j]*h[0][j]);
    rh[m].chi[j] = creal(1.I*beta_p[j]*h[0][j] - h[1][j]/zeta_p[j]);
  }
}


void rk4(params_ptr in) {
  tmp = in->w_ptr;
  memcpy(aux.rho, tmp->rho, N*sizeof(fftw_complex));
  memcpy(aux.chi, tmp->chi, N*sizeof(fftw_complex)); 
  aux.alpha = tmp->alpha;
  rhs(tmp,0);
  for (int j = 0; j < N; j++) {
    tmp->alpha  = aux.alpha  + 0.5*dt*rh[0].alpha;
    tmp->rho[j] = aux.rho[j] + 0.5*dt*rh[0].rho[j];
    tmp->chi[j] = aux.chi[j] + 0.5*dt*rh[0].chi[j];
  }
  rhs(tmp, 1);
  for (int j = 0; j < N; j++) {
    tmp->alpha  = aux.alpha  + 0.5*dt*rh[1].alpha;
    tmp->rho[j] = aux.rho[j] + 0.5*dt*rh[1].rho[j];
    tmp->chi[j] = aux.chi[j] + 0.5*dt*rh[1].chi[j];
  }
  rhs(tmp, 2);
  for (int j = 0; j < N; j++) {
    tmp->alpha  = aux.alpha  + 0.5*dt*rh[2].alpha;
    tmp->rho[j] = aux.rho[j] + dt*rh[2].rho[j];
    tmp->chi[j] = aux.chi[j] + dt*rh[2].chi[j];
  }
  rhs(tmp, 3);
  tmp->alpha  = aux.alpha + (1./6)*dt*(rh[0].alpha + 2.*rh[1].alpha + 2.*rh[2].alpha + rh[3].alpha);
  for (int j = 0; j < N; j++) {
    tmp->rho[j] = aux.rho[j] + (1./6)*dt*(rh[0].rho[j] + 2.*rh[1].rho[j] + 2.*rh[2].rho[j] + rh[3].rho[j]);
    tmp->chi[j] = aux.chi[j] + (1./6)*dt*(rh[0].chi[j] + 2.*rh[1].chi[j] + 2.*rh[2].chi[j] + rh[3].chi[j]);
  }
  //fforcing(in);
  ffilter(in);
}
/*
void generate_derivatives(work_ptr in, int m) {
  char filename[80];
  sprintf(filename, "der_%02d.dat", 0);
  reconstruct_surface(in);
  write_out(in->Z, filename);
  for (int j = 1; j < m + 1; j ++) {
    der(in->Z, in->Z);
    sprintf(filename, "der_%02d.dat", j);
    write_out(in->Z, filename);
  }
}
*/
























